#include <algorithm/analytical_model/CommonPara.h>
#include <algorithm/analytical_model/HeliostatPara.h>
#include <algorithm/analytical_model/HeliostatWithGlassPara.h>
#include <algorithm/analytical_model/ScenePara.h>
#include <algorithm/analytical_model/algorithm/ModelPara.h>
#include <algorithm/analytical_model/algorithm/dnn/backend/openvino/CompiledModelCache.h>
#include <algorithm/analytical_model/algorithm/dnn/params/UNIZAR2.h>
#include <algorithm/analytical_model/algorithm/geometry/AllGeometry.h>
#include <algorithm/analytical_model/algorithm/geometry/Utils.h>
#include <algorithm/analytical_model/convolution/UNIZARPara2.h>
#include <algorithm/shared/data/scene/Heliostat.h>
#include <algorithm/shared/data/scene/Receiver.h>
#include <algorithm/shared/data/scene/Sun.h>
#include <algorithm/shared/data/topology/geometry/Rectangle3D.h>
#include <algorithm/shared/data/world/sun/sun_position/SunPosition.h>
#include <utils/ParaUtils.h>
#include <utils/config.h>

#include <openvino/openvino.hpp>

#include <bit>
#include <memory>
#include <vector>

// 仅支持矩形平面定日镜

using namespace solar;

template <>
void solar::computeSigma(SimulationPara* simulation_para, ScenePara* scene_para,
                         SOLAR_ARRAY_ARG HeliostatPara* heliostat_para, CommonPara* common_para,
                         Sun* sun, SOLAR_ARRAY_ARG Heliostat* heliostat, Receiver* receiver,
                         std::size_t num_helios, SOLAR_ARRAY_ARG UNIZARPara2* ptr,
                         ComputeSigmaCallback callback)
{
    std::vector<dnn::UNIZAR2Input> net_inputs(num_helios);
    const auto model_path = SOLAR_GET_VALUE(CommonPara, path_to_model_, *common_para);
    auto compiled_model_cache = dnn::openvino::CompiledModelCache::get_mutable_instance();
    dnn::openvino::DeviceType preferred_device_type =
        num_helios < 100 ? dnn::openvino::DeviceType::CPU : dnn::openvino::DeviceType::GPU;
    ov::CompiledModel& compiled_model =
        compiled_model_cache.getCompiledModel(model_path, preferred_device_type);
    constexpr ov::element::Type input_type = ov::element::f32;
    constexpr auto input_params_count = sizeof(dnn::UNIZAR2Input) / sizeof(float);
    // constexpr auto output_params_count = sizeof(dnn::UNIZAR2Output) / sizeof(float);

    constexpr auto original_reflectivity = 0.96;
    constexpr auto original_dni = 1000.0;
    const auto sun_position_ptr = SOLAR_GET_VALUE(Sun, sun_position_, *sun);
    const auto azimuth = SOLAR_GET_VALUE(SunPosition, azimuth_, *sun_position_ptr);
    const auto elevation = SOLAR_GET_VALUE(SunPosition, elevation_, *sun_position_ptr);
    const auto dni = SOLAR_GET_VALUE(ScenePara, DNI_, *scene_para);
    const auto csr = SOLAR_GET_VALUE(ScenePara, CSR_, *scene_para);

    SOLAR_OMP_PARALLEL_FOR_SIMD
    for (std::size_t index = 0; index < num_helios; index++)
    {
        const auto geometry_heliostat = SOLAR_GET_VALUE(Heliostat, geometry_, heliostat[index]);
        const auto geometry_receiver = SOLAR_GET_VALUE(Receiver, geometry_, *receiver);
        const auto sigma_slope_error =
            SOLAR_GET_VALUE(HeliostatPara, sigma_slope_error_, heliostat_para[index]);

        const auto heliostat_center = ::solar::getCenter(geometry_heliostat.get());
        const auto focus_center =
            ::solar::getFocusCenter(geometry_heliostat.get(), geometry_receiver.get());
        const auto distance_heliostat_focus_center = (heliostat_center - focus_center).length();
        const auto heliostat_area = solar::getArea(geometry_heliostat.get());
        const auto geometry_heliostat_rectangle3d =
            std::dynamic_pointer_cast<Rectangle3D>(geometry_heliostat);
        const auto geometry_heliostat_rectangle3d_size =
            SOLAR_GET_VALUE(Rectangle3D, size_, *geometry_heliostat_rectangle3d);
        const auto heliostat_l_w_ratio =
            geometry_heliostat_rectangle3d_size.x() / geometry_heliostat_rectangle3d_size.z();
        const auto heliostat_glass_para =
            SOLAR_GET_VALUE(HeliostatPara, heliostat_glass_para_, heliostat_para[index]);
        const auto heliostat_glass_thickness =
            SOLAR_GET_VALUE(HeliostatWithGlassPara, glass_thickness_, *heliostat_glass_para);
        const auto heliostat_glass_refractive_index =
            SOLAR_GET_VALUE(HeliostatWithGlassPara, refractivity_, *heliostat_glass_para);
        net_inputs[index] = {
            .heliostat_area_ = static_cast<float>(heliostat_area),
            .heliostat_l_w_ratio_ = static_cast<float>(heliostat_l_w_ratio),
            .distance_heliostat_receiver_ = static_cast<float>(distance_heliostat_focus_center),
            .heliostat_glass_thickness_ =
                static_cast<float>(heliostat_glass_thickness * 1000.0), // 转换为毫米
            .sun_elevation_ = static_cast<float>(elevation),
            .sun_azimuth_ = static_cast<float>(azimuth),
            .heliostat_glass_sigma_ = static_cast<float>(sigma_slope_error.x()),
            .sun_csr_ = static_cast<float>(csr),
            .heliostat_glass_refractive_index_ =
                static_cast<float>(heliostat_glass_refractive_index)};
    }

    auto infer_request = compiled_model.create_infer_request();
    auto* input_data = std::bit_cast<float*>(net_inputs.data());
    ov::Shape input_shape = {num_helios, input_params_count};
    ov::Tensor input_tensor = {input_type, input_shape, input_data};
    infer_request.set_input_tensor(input_tensor);
    infer_request.infer();
    const auto& output_tensor = infer_request.get_output_tensor();
    auto* output_ptr = std::bit_cast<dnn::UNIZAR2Output*>(output_tensor.data<float>());

    const auto dni_ratio = dni / original_dni;
    SOLAR_OMP_PARALLEL_FOR_SIMD
    for (std::size_t index = 0; index < num_helios; index++)
    {
        const auto reflectivity =
            SOLAR_GET_VALUE(HeliostatPara, reflectivity_, heliostat_para[index]);
        const auto reflectivity_ratio = reflectivity / original_reflectivity;
        output_ptr[index].peak_flux_param_ *= dni_ratio * reflectivity_ratio;
    }
    // Run callback function
    callback(output_ptr);

    SOLAR_OMP_PARALLEL_FOR_SIMD
    for (std::size_t index = 0; index < num_helios; index++)
    {
        UNIZARPara2 current_param;
        SOLAR_SET_VALUE(UNIZARPara2, flux_param_, output_ptr[index].peak_flux_param_, current_param)
        SOLAR_SET_VALUE(UNIZARPara2, sigma_x_, output_ptr[index].sigma_x_, current_param)
        SOLAR_SET_VALUE(UNIZARPara2, sigma_y_, output_ptr[index].sigma_y_, current_param)
        SOLAR_SET_VALUE(UNIZARPara2, x3_, output_ptr[index].x3_, current_param)
        SOLAR_SET_VALUE(UNIZARPara2, x0_, output_ptr[index].x0_, current_param)
        SOLAR_SET_VALUE(UNIZARPara2, y1_, output_ptr[index].y1_, current_param)
        SOLAR_SET_VALUE(UNIZARPara2, y0_, output_ptr[index].y0_, current_param)
        ptr[index] = current_param;
    }
}
